{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import cPickle as pickle\n",
    "from IPython.display import display\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.backends.backend_pdf import PdfPages\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "remap_dict = {\n",
    "    ('unmeasure', 'blur'): 'unmeasure-blur',\n",
    "    ('unmeasure', 'inpaint-tv') : 'unmeasure-inpaint-tv',\n",
    "    ('baseline', 'None') : 'ignore',\n",
    "    ('ambient', 'None') : 'AmbientGAN (ours)',\n",
    "    ('unmeasure', 'wiener'): 'unmeasure-weiner',\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def convert_df(df, convert_spec):\n",
    "    for colname in convert_spec['float']:\n",
    "        df.loc[:, colname] = df[colname].apply(np.float32)\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def filter_df(df, filter_spec):\n",
    "    df = df[filter_spec['cols']]\n",
    "    filter_ = True\n",
    "    for colname, values in filter_spec['in'].iteritems():\n",
    "        filter_ = filter_ & df[colname].isin(values)\n",
    "    df = df[filter_]\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def save_plot(save_path):\n",
    "    pdf = PdfPages(save_path)\n",
    "    pdf.savefig(bbox_inches='tight')\n",
    "    pdf.close()    \n",
    "\n",
    "def errorbar(x, y, std, **kwargs):\n",
    "    (_, caps, _) = plt.errorbar(x, y, yerr=1.96*std,\n",
    "                                marker='o', markersize=5, capsize=5, linewidth=2, **kwargs)\n",
    "    for cap in caps:\n",
    "        cap.set_markeredgewidth(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def plot1(plot_spec, df):\n",
    "\n",
    "    figure_groups = df.groupby(plot_spec['figure_group_by'])\n",
    "    for fig_key, fig_val in figure_groups.groups.iteritems():\n",
    "        fig_group_df = df.loc[fig_val, :]\n",
    "\n",
    "        plt.figure(figsize=[6, 4])\n",
    "        # plt.title(fig_key)\n",
    "        \n",
    "        curve_groups = fig_group_df.groupby(plot_spec['curve_group_by'])\n",
    "        x_var = plot_spec['x_var']\n",
    "        y_var = plot_spec['y_var']\n",
    "        std_var = plot_spec['std_var']\n",
    "        \n",
    "        legends = []\n",
    "\n",
    "        for curve_key in sorted(curve_groups.groups):\n",
    "            curve_val = curve_groups.groups[curve_key]\n",
    "            # for curve_key, curve_val in curve_groups.groups.iteritems():\n",
    "            group_df = fig_group_df.loc[curve_val, :]\n",
    "            legends.append(remap_dict[curve_key])\n",
    "\n",
    "            cols = [x_var, y_var, std_var]\n",
    "            group_df.sort_values(by=[x_var], inplace=True)\n",
    "            x = np.array(group_df[x_var])\n",
    "            y = np.array(group_df[y_var])\n",
    "\n",
    "            if std_var is not None:\n",
    "                std = np.array(group_df[std_var])\n",
    "                errorbar(x, y, std, alpha=0.8)\n",
    "            else:\n",
    "                plt.plot(x, y, '-o')\n",
    "\n",
    "        ## Prettify\n",
    "        # axis\n",
    "        plt.gca().set_ylim(bottom=0)\n",
    "        # plt.gca().set_xscale(\"log\", nonposx='clip')\n",
    "        # plt.gca().set_xlim([-0.1, 1.1])\n",
    "\n",
    "        # labels, ticks, titles\n",
    "        # ticks = [10, 25, 50, 100, 200, 300, 400, 500, 750]\n",
    "        # labels = [10, 25, 50, 100, 200, 300, 400, 500, 750]\n",
    "        # plt.xticks(ticks, labels, rotation=90)\n",
    "        plt.xlabel(plot_spec['x_label'], fontsize=14)\n",
    "        plt.ylabel(plot_spec['y_label'], fontsize=14)\n",
    "\n",
    "        # Legends\n",
    "        plt.legend(legends, fontsize=12.5, loc=plot_spec['legend_loc'])\n",
    "        plt.ylim(plot_spec['y_lim'])\n",
    "\n",
    "        if plot_spec['save']:\n",
    "            figname = '_'.join(fig_key)\n",
    "            fig_savepath = './results/plots/{}.pdf'.format(figname)\n",
    "            save_plot(fig_savepath)\n",
    "               \n",
    "def plot2(df):\n",
    "    fig = plt.figure(figsize=[6, 4])\n",
    "    ax = plt.subplot(111)\n",
    "    legends = []\n",
    "    legend_dict = {'ambient': 'AmbientGAN', 'baseline': 'ignore', 'unmeasure': 'unmeasure-blur'}\n",
    "    colors = {'ambient': 'blue', 'baseline': 'green', 'unmeasure': 'red'}\n",
    "    df.sort_values(by=['train_mode', 'drop_prob'], inplace=True)\n",
    "    for i, row in df.iterrows():\n",
    "        drop_prob = np.around(row['drop_prob'], decimals=2)\n",
    "        train_mode = row['train_mode']\n",
    "        legend = 'p = {}, {}'.format(legend_dict[train_mode], drop_prob)\n",
    "        legends.append(legend)\n",
    "        iters, scores = zip(*row['m_inception'])\n",
    "        mean, std = map(np.array, zip(*scores))\n",
    "        errorbar(iters, mean, std, alpha=1-drop_prob/1.1, color=colors[train_mode])\n",
    "    # plt.legend(legends, fontsize=12.5, loc=2)\n",
    "    plt.xlabel('Training iteration',  fontsize=14)\n",
    "    plt.ylabel('Inception score', fontsize=14)\n",
    "    \n",
    "    #     # Following block is based on https://stackoverflow.com/a/4701285/3537687\n",
    "    #     # Shrink current axis by 20%\n",
    "    #     box = ax.get_position()\n",
    "    #     ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])\n",
    "    #     # Put a legend to the right of the current axis\n",
    "    #     ax.legend(legends, loc='center left', bbox_to_anchor=(1, 0.5))\n",
    "    \n",
    "    plt.ylim([0, 8.2])\n",
    "    \n",
    "    fig_savepath = './results/plots/cifar10_drop_independent_vs-time.pdf'\n",
    "    save_plot(fig_savepath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def produce_plot1(df_filepath, convert_spec, filter_spec, plot_spec):\n",
    "    with open(df_filepath, 'rb') as df_pkl_file:\n",
    "        df = pickle.load(df_pkl_file)\n",
    "    \n",
    "    df = convert_df(df, convert_spec)    \n",
    "    df = filter_df(df, filter_spec)    \n",
    "    display(df)\n",
    "    plot1(plot_spec, df)\n",
    "    \n",
    "def produce_plot2(df_filepath, convert_spec, filter_spec):\n",
    "    with open(df_filepath, 'rb') as df_pkl_file:\n",
    "        df = pickle.load(df_pkl_file)\n",
    "    df = convert_df(df, convert_spec)    \n",
    "    df = filter_df(df, filter_spec)    \n",
    "    #     display(df)\n",
    "    plot2(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def main1():\n",
    "    df_filepath = './results/df_cifar.pkl'\n",
    "\n",
    "    convert_spec = {\n",
    "        'float': [\n",
    "            'drop_prob',\n",
    "        ]        \n",
    "    }\n",
    "\n",
    "    filter_spec = {\n",
    "\n",
    "        'cols' : [\n",
    "            'dataset', \n",
    "            'measurement_type', \n",
    "            'train_mode', \n",
    "            'drop_prob',\n",
    "            'unmeasure_type', \n",
    "            # 'm_inception',\n",
    "            'm_final_score',\n",
    "            'm_final_std',\n",
    "        ],\n",
    "\n",
    "        'in': {\n",
    "            'measurement_type': [\n",
    "                'drop_independent'\n",
    "            ]\n",
    "        }\n",
    "\n",
    "    }\n",
    "\n",
    "    plot_spec = {\n",
    "        'figure_group_by' : ['dataset', 'measurement_type'],    \n",
    "        'curve_group_by': ['train_mode', 'unmeasure_type'],\n",
    "        # change these to lambda functions if needed\n",
    "        'x_var': 'drop_prob',\n",
    "        'y_var': 'm_final_score',\n",
    "        'std_var': 'm_final_std',\n",
    "        'x_label': 'Block probability (p)',\n",
    "        'y_label': 'Inception score',\n",
    "        'save': True,\n",
    "        'y_lim': [0, 8.2],\n",
    "        'legend_loc': 3,\n",
    "    }\n",
    "\n",
    "    produce_plot1(df_filepath, convert_spec, filter_spec, plot_spec)\n",
    "    \n",
    "main1()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main2():\n",
    "    df_filepath = './results/df_cifar.pkl'\n",
    "\n",
    "    convert_spec = {\n",
    "        'float': [\n",
    "            'drop_prob',\n",
    "        ]        \n",
    "    }\n",
    "\n",
    "    filter_spec = {\n",
    "\n",
    "        'cols' : [\n",
    "            'dataset', \n",
    "            'measurement_type', \n",
    "            'train_mode', \n",
    "            'drop_prob',\n",
    "            'unmeasure_type', \n",
    "            'm_inception',\n",
    "        ],\n",
    "\n",
    "        'in': {\n",
    "            'measurement_type': [\n",
    "                'drop_independent'\n",
    "            ],\n",
    "#             'train_mode': ['ambient', 'baseline']\n",
    "        }\n",
    "\n",
    "    }\n",
    "    \n",
    "    produce_plot2(df_filepath, convert_spec, filter_spec)\n",
    "    \n",
    "main2()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# def main3():\n",
    "#     df_filepath = './results/df_cifar10.pkl'\n",
    "\n",
    "#     convert_spec = {\n",
    "#         'float': [\n",
    "#             'drop_prob',\n",
    "#             'additive_noise_std',\n",
    "#         ]        \n",
    "#     }\n",
    "\n",
    "#     filter_spec = {\n",
    "\n",
    "#         'cols' : [\n",
    "#             'dataset', \n",
    "#             'measurement_type', \n",
    "#             'train_mode', \n",
    "#             'additive_noise_std',\n",
    "#             'unmeasure_type', \n",
    "#             # 'm_inception',\n",
    "#             'm_final_score',\n",
    "#             'm_final_std',\n",
    "#         ],\n",
    "\n",
    "#         'in': {\n",
    "#             'measurement_type': [\n",
    "#                 'blur_addnoise'\n",
    "#             ]\n",
    "#         }\n",
    "\n",
    "#     }\n",
    "\n",
    "#     plot_spec = {\n",
    "#         'figure_group_by' : ['dataset', 'measurement_type'],    \n",
    "#         'curve_group_by': ['train_mode', 'unmeasure_type'],\n",
    "#         # change these to lambda functions if needed\n",
    "#         'x_var': 'additive_noise_std',\n",
    "#         'y_var': 'm_final_score',\n",
    "#         'std_var': 'm_final_std',\n",
    "#         'x_label': 'Noise standard deviation ($\\sigma$)',\n",
    "#         'y_label': 'Inception score',\n",
    "#         'save': True,\n",
    "#         'y_lim': [0, None],\n",
    "#         'legend_loc': 1,\n",
    "#     }\n",
    "\n",
    "#     produce_plot1(df_filepath, convert_spec, filter_spec, plot_spec)\n",
    "    \n",
    "# main3()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
